package pwd.allen.plugin;

import cn.hutool.core.date.DateUtil;
import cn.hutool.core.util.StrUtil;
import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.SystemMetaObject;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.sql.Connection;
import java.text.ParseException;
import java.text.SimpleDateFormat;
import java.util.*;
import java.util.regex.Pattern;

@Intercepts({
        @Signature(type = StatementHandler.class, method = "prepare", args = {Connection.class, Integer.class}),
        @Signature(type = Executor.class, method = "update", args = {MappedStatement.class, Object.class}),
        @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}),
        @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class, CacheKey.class, BoundSql.class})
})
public class ShardDatePlugin implements Interceptor {
    private static final Logger logger = LoggerFactory.getLogger(ShardDatePlugin.class);

    public static final SimpleDateFormat DATE_FORMAT_yyyyMM = new SimpleDateFormat(ShardDateInfo.SUB_MODE_YYYYMM);
    public static final SimpleDateFormat DATE_FORMAT_yyyy = new SimpleDateFormat(ShardDateInfo.SUB_MODE_YYYY);

    /**
     * 业务参数存储，用于决定分表策略
     */
    public static final ThreadLocal<ShardDateParam> THREAD_LOCAL = new ThreadLocal<>();

    /**
     * 表名映射
     */
    private final Map<String, Map<String, ShardDateInfo>> tableMap = new HashMap<>();

    /**
     * @param invocation
     * @return
     * @throws Throwable
     */
    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        try {
            if ("update".equals(invocation.getMethod().getName())) {
                return update(invocation);
            } else if ("query".equals(invocation.getMethod().getName())) {
                return query(invocation);
            } else if ("prepare".equals(invocation.getMethod().getName())) {
                return prepare(invocation);
            }
            return invocation.proceed();
        } catch (Throwable throwable) {
            logger.error(throwable.toString(), throwable);
            throw throwable;
        }
    }

    private Object prepare(Invocation invocation) throws InvocationTargetException, IllegalAccessException {
        ShardDateParam shardDateParam = THREAD_LOCAL.get();
        try {
            if (shardDateParam != null && shardDateParam.getNewSql() != null) {
                BoundSql boundSql = StatementHandler.class.cast(invocation.getTarget()).getBoundSql();
                MetaObject metaObjectBoundSql = SystemMetaObject.forObject(boundSql);
                metaObjectBoundSql.setValue("sql", shardDateParam.getNewSql());
            }
            return invocation.proceed();
        } finally {
            if (shardDateParam != null) {
                if (shardDateParam.isIfRemoveAfterPrepare()) {
                    THREAD_LOCAL.remove();
                } else {
                    shardDateParam.setIfRemoveAfterPrepare(true);
                }
            }
        }
    }

    /**
     * 更新、删除、插入分表处理
     * @param invocation
     * @return
     * @throws Throwable
     */
    private Object update(Invocation invocation) throws Throwable {
        ShardDateParam shardDateParam = THREAD_LOCAL.get();

        if (shardDateParam == null) {
//            shardDateParam = new ShardDateParam(DATE_FORMAT_yyyyMM, null, null, null);
//            THREAD_LOCAL.set(shardDateParam);
            return invocation.proceed();
        }

        Executor executor = (Executor) invocation.getTarget();
        Object[] args = invocation.getArgs();
        MappedStatement ms = (MappedStatement) args[0];
        Object parameterObject = args[1];
        BoundSql boundSql = ms.getBoundSql(parameterObject);
        String sql = boundSql.getSql();

        //<editor-fold desc="校验是否需要分表处理">
        String sqlUpper = sql.trim().toUpperCase();
        for (String tableName : getTableMap().keySet()) {
            if (sqlUpper.matches("[\\s\\S]*\\s" + tableName + "(\\s|\\()[\\s\\S]*")) {
                shardDateParam.setTableName(tableName);
                break;
            }
        }
        if (shardDateParam.getSimpleDateFormat() == null) {
            throw new RuntimeException(String.format("更新时分表配置需要指定分表模式！"));
        }
        if (shardDateParam.getTableName() == null) {
            throw new RuntimeException(String.format("sql【%s】中匹配不到分表", ms.getId()));
        }
        //</editor-fold>

        // 分表策略
        ShardDateInfo shardDateInfo = getTableMap(shardDateParam.getTableName(), shardDateParam.getSimpleDateFormat().toPattern());
        if (shardDateInfo == null) {
            throw new RuntimeException(String.format("%s分表没有配置指定的模式：%s", shardDateParam.getTableName(), shardDateParam.getSimpleDateFormat().toPattern()));
        }

        SimpleDateFormat dateFormat = null;
        if (ShardDateInfo.SUB_MODE_YYYY.equals(shardDateInfo.getSubMode())) {
            dateFormat = DATE_FORMAT_yyyy;
        } else if (ShardDateInfo.SUB_MODE_YYYYMM.equals(shardDateInfo.getSubMode())) {
            dateFormat = DATE_FORMAT_yyyyMM;
        }

        //<editor-fold desc="按照日期存储">
        HashMap<String, List> dateMap = new HashMap<>();
        List listBatch = null;
        if (parameterObject instanceof List) {
            listBatch = (List) parameterObject;
        } else if (parameterObject instanceof Map && Map.class.cast(parameterObject).containsKey("list")) {
            listBatch = (List) Map.class.cast(parameterObject).get("list");
        } else {
            listBatch = new ArrayList();
            listBatch.add(parameterObject);
        }
        if (listBatch.size() == 0) {
            return 0;
        }
        Method method = getMethod(listBatch.get(0), shardDateInfo.getFieldName());
        if (method == null) {
            throw new IllegalArgumentException(String.format("%s分表策略配置有误，%s字段找不到", shardDateParam.getTableName(), shardDateInfo.getFieldName()));
        }
        for (Object data : listBatch) {
            String dateStr = null;
            Date valDate = null;
            Object val = method.invoke(data, null);
            if (val == null) {
                throw new IllegalArgumentException(String.format("%s分表字段%s不能为空", shardDateInfo.getFieldName()));
            }
            if (val instanceof Date) {
                valDate = Date.class.cast(val);

            } else {
                valDate = dateFormat.parse(val.toString());
            }
            if (valDate.before(shardDateInfo.getInitDate())) {
                valDate = shardDateInfo.getInitDate();
            }
            dateStr = dateFormat.format(valDate);
            List dataList = dateMap.get(dateStr);
            if (dataList == null) {
                dataList = new ArrayList<>();
                dateMap.put(dateStr, dataList);
            }
            dataList.add(data);
        }
        //</editor-fold>

        int rel = 0;
        Pattern pattern = Pattern.compile(shardDateParam.getTableName(), Pattern.CASE_INSENSITIVE);

        Iterator<Map.Entry<String, List>> iterator = dateMap.entrySet().iterator();
        if (dateMap.size() == 1) {
            // 只有一个日期时，不需要拆分list，直接调用即可
            Map.Entry<String, List> entry = iterator.next();
            sql = pattern.matcher(sql).replaceAll(shardDateInfo.getNewTablePrefix() + "_" + entry.getKey());
            shardDateParam.setNewSql(sql);
            Object proceed = invocation.proceed();
            return proceed;
        } else {
            // 如果日期有多个，则拆分了list，需要重新组装sql
            while (iterator.hasNext()) {
                Map.Entry<String, List> entry = iterator.next();
                List list = entry.getValue();
                if (parameterObject instanceof List) {
                    parameterObject = list;
                } else if (parameterObject instanceof Map) {
                    Map.class.cast(parameterObject).put("list", list);
                }
                boundSql = ms.getBoundSql(parameterObject);
                sql = pattern.matcher(boundSql.getSql()).replaceAll(shardDateInfo.getNewTablePrefix() + "_" + entry.getKey());
                shardDateParam.setNewSql(sql);
                if (iterator.hasNext()) {
                    shardDateParam.setIfRemoveAfterPrepare(false);
                }
                int update = executor.update(ms, parameterObject);
                rel += update;
            }
            return rel;
        }
    }

    /**
     * 查询分表处理
     * @param invocation
     * @return
     * @throws Throwable
     */
    private Object query(Invocation invocation) throws Throwable {
        ShardDateParam shardDateParam = THREAD_LOCAL.get();

        Object[] args = invocation.getArgs();
        MappedStatement ms = (MappedStatement) args[0];
        Object parameter = args[1];

        BoundSql boundSql;
        //由于逻辑关系，只会进入一次
        if (args.length == 4) {
            //4 个参数时
            boundSql = ms.getBoundSql(parameter);
        } else {
            //6 个参数时
            boundSql = (BoundSql) args[5];
        }

        String sql = boundSql.getSql();

        //<editor-fold desc="校验是否需要分表处理">
        if (shardDateParam == null) {
            return invocation.proceed();
        }
        if ("com.cdc.common.orm.ShardDateInfo.selectAll".equals(ms.getId())) {
            // 获取分表配置的查询需要放行
            shardDateParam.setIfRemoveAfterPrepare(false);
            return invocation.proceed();
        }
        if (shardDateParam.getBeginDate() == null || shardDateParam.getEndDate() == null) {
            throw new RuntimeException("分表的起始时间和结束时间不能为空");
        }
        if (shardDateParam.getBeginDate().after(shardDateParam.getEndDate())) {
            throw new RuntimeException("起始时间和结束时间不符合");
        }
        if (shardDateParam.getSimpleDateFormat() == null) {
            // 默认按照yyyyMM分表
            shardDateParam.setSimpleDateFormat(DATE_FORMAT_yyyyMM);
        }
        if (shardDateParam.getTableName() == null) {
            // 如果分表配置没有指定表名，则根据sql自动匹配，匹配到多个的话通过链表形势链接，只支持同一种分表模式
            String sqlUpper = sql.trim().toUpperCase();
            for (String tableName : getTableMap().keySet()) {
                if (sqlUpper.matches("[\\s\\S]*\\s" + tableName + "(\\s|\\()[\\s\\S]*")) {
                    ShardDateParam param2Set = shardDateParam;
                    while (param2Set.getNext() != null) {
                        param2Set = param2Set.getNext();
                    }
                    if (StrUtil.isEmpty(param2Set.getTableName())) {
                        param2Set.setTableName(tableName);
                    } else {
                        param2Set.setNext(param2Set.deepClone());
                        param2Set.getNext().setTableName(tableName);
                    }
                }
            }
            if (shardDateParam.getTableName() == null) {
                throw new RuntimeException(String.format("sql【%s】中匹配不到分表", ms.getId()));
            }
        }
        //</editor-fold>

        //<editor-fold desc="分表策略">
        for (ShardDateParam param2Set = shardDateParam; param2Set != null; param2Set = param2Set.getNext()) {
            ShardDateInfo shardDateInfo = getTableMap(param2Set.getTableName(), param2Set.getSimpleDateFormat().toPattern());
            if (shardDateInfo == null) {
                throw new RuntimeException(String.format("%s分表没有配置指定的模式：%s", param2Set.getTableName(), param2Set.getSimpleDateFormat().toPattern()));
            }
            param2Set.setShardDateInfo(shardDateInfo);
            //<editor-fold desc="记录所有分表的日期">
            List<String> listDate = new ArrayList<>();
            Calendar calendar = Calendar.getInstance();
            calendar.setTime(shardDateInfo.getInitDate().after(param2Set.getBeginDate()) ? shardDateInfo.getInitDate() : param2Set.getBeginDate());
            while (!calendar.getTime().after(param2Set.getEndDate())) {
                listDate.add(param2Set.getSimpleDateFormat().format(calendar.getTime()));
                if (ShardDateInfo.SUB_MODE_YYYYMM.equals(shardDateInfo.getSubMode())) {
                    calendar.add(Calendar.MONTH, 1);
                } else if (ShardDateInfo.SUB_MODE_YYYY.equals(shardDateInfo.getSubMode())) {
                    calendar.add(Calendar.YEAR, 1);
                }
            }
            if (listDate.size() == 0) {
                throw new RuntimeException(String.format("%s分表的%s模式没有匹配的日期", shardDateInfo.getTableName(), shardDateInfo.getSubMode()));
            }
            param2Set.setDateStrs(listDate);
            //</editor-fold>
        }
        //</editor-fold>


        //<editor-fold desc="拼接sql">
        StringBuilder sb = new StringBuilder();
        String[] split = sql.toUpperCase().split("\\)\\s*SELECT\\s");
        if (split.length > 1) {
            // sql中已经存在with ... as ()这种写法，则将with 的内容先拼接进去
            sb.append(split[0]).append(")");
            sql = sql.substring(split[0].length() + 1);
        }
        for (ShardDateParam param2Set = shardDateParam; param2Set != null; param2Set = param2Set.getNext()) {
            if (sb.length() == 0) {
                sb.append("with ");
            } else {
                sb.append(",");
            }
            sb.append(param2Set.getTableName()).append(" as\n(\n");
            for (int i = 0; i < param2Set.getDateStrs().size(); i++) {
                if (i != 0) {
                    sb.append("\tunion all\n");
                }
                sb.append(String.format("\tselect * from %s_%s\n", param2Set.getShardDateInfo().getNewTablePrefix(), param2Set.getDateStrs().get(i)));
            }
            sb.append(")\n");
        }
        sql = sb.toString() + sql;
        //</editor-fold>

        //修改sql
        shardDateParam.setNewSql(sql);

        //执行目标方法
//        List<Object> query = executor.query(ms, parameter, rowBounds, resultHandler, cacheKey, boundSql);
//        return query;
        return invocation.proceed();
    }

    @Override
    public Object plugin(Object target) {
        return Plugin.wrap(target, this);
    }



    private Map<String, Map<String, ShardDateInfo>> getTableMap() {
        if (tableMap.size() == 0) {
            // 初始化分表配置
            List<ShardDateInfo> listConfig = new ArrayList<>();
            listConfig.add(new ShardDateInfo("DB_USER", "DB_USER", "createAt", ShardDateInfo.SUB_MODE_YYYYMM, 1, DateUtil.parse("2022-07", "yyyy-MM")));
            for (ShardDateInfo shardDateInfo : listConfig) {
                setInitDateIfNotExist(shardDateInfo);
                Map<String, ShardDateInfo> stringShardDateInfoMap = tableMap.get(shardDateInfo.getTableName());
                if (stringShardDateInfoMap == null) {
                    stringShardDateInfoMap = new HashMap<>();
                    tableMap.put(shardDateInfo.getTableName(), stringShardDateInfoMap);
                }
                stringShardDateInfoMap.put(shardDateInfo.getSubMode(), shardDateInfo);
            }
        }
        return tableMap;
    }
    private ShardDateInfo getTableMap(String tableName, String format) {
        Map<String, ShardDateInfo> stringShardDateInfoMap = getTableMap().get(tableName);
        if (stringShardDateInfoMap != null) {
            return stringShardDateInfoMap.get(format);
        }
        return null;
    }

    public static Method getMethod(Object obj, String fieldName) {
        Class clazz = obj.getClass();
        String methodName = "get" + fieldName.substring(0, 1).toUpperCase() + fieldName.substring(1);
        Method method = null;
        while (method == null && clazz != Object.class) {
            try {
                method = clazz.getDeclaredMethod(methodName);
            } catch (Exception e) {
                clazz = clazz.getSuperclass();
            }
        }
        return method;
    }

    public static void setInitDateIfNotExist(ShardDateInfo shardDateInfo) {
        try {
            if (shardDateInfo != null && shardDateInfo.getInitDate() == null && shardDateInfo.getDeleteNumber() != null && shardDateInfo.getSubMode() != null) {
                Calendar calendar = Calendar.getInstance();
                Date dateNow = new Date();
                if (ShardDateInfo.SUB_MODE_YYYY.equals(shardDateInfo.getSubMode())) {
                    calendar.setTime(DATE_FORMAT_yyyy.parse(DATE_FORMAT_yyyy.format(dateNow)));
                    calendar.add(Calendar.YEAR, -shardDateInfo.getDeleteNumber());
                } else if (ShardDateInfo.SUB_MODE_YYYYMM.equals(shardDateInfo.getSubMode())) {
                    calendar.setTime(DATE_FORMAT_yyyyMM.parse(DATE_FORMAT_yyyyMM.format(dateNow)));
                    calendar.add(Calendar.MONTH, -shardDateInfo.getDeleteNumber());
                }
                shardDateInfo.setInitDate(calendar.getTime());
            }
        } catch (ParseException e) {}
    }

    @Override
    public void setProperties(Properties properties) {
    }
}





